package com.etocrm.saas.lib.base.resource.config;

import com.etocrm.saas.lib.base.common.result.Result;
import com.etocrm.saas.lib.base.resource.Util.IpAddressUtil;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 令牌桶算法限流控制
 */
@Slf4j
@Aspect
@Component
public class LimitAopAspect {


    private final ConcurrentHashMap<String, RateLimiter> limitMap = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, RateLimiter> limitAllMap = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, RateLimiter> limitIpMap = new ConcurrentHashMap<>();

    @Pointcut("@annotation(limit)")
    public void excudeService(Limit limit) {
    }

    @Around("excudeService(limit)")
    public Object around(ProceedingJoinPoint joinPoint,Limit limit) throws Throwable{
        RequestAttributes ra = RequestContextHolder.getRequestAttributes();
        ServletRequestAttributes sra = (ServletRequestAttributes) ra;
        HttpServletRequest request = sra.getRequest();
        if(limit.isAppointIp()){
            return appointIpCount(joinPoint,limit,request);
        }
        //key作用：不同的接口，不同的流量控制
        String key = request.getRequestURI();
        RateLimiter rateLimiter = null;
        //验证缓存是否有命中key
        if (!limitMap.containsKey(key)) {
            //创建令牌桶,指定令牌数量
            rateLimiter = RateLimiter.create(limit.permitsPerSecond());
            limitMap.put(key, rateLimiter);
            log.info("新建了令牌桶={}，容量={}",key,limit.permitsPerSecond());
        }
        rateLimiter = limitMap.get(key);
        //拿令牌
        boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
        //拿不到命令，直接返回异常提示
        if (!acquire) {
            log.info("令牌桶={}，获取令牌失败",key);
            return Result.fail(limit.msg());
        }
        log.info("令牌桶={}，获取令牌成功",key);
        return joinPoint.proceed();
    }

    private Object appointIpCount(ProceedingJoinPoint joinPoint,Limit limit,HttpServletRequest request)throws Throwable{
        double ipCount = limit.ipCount();
        if(ipCount <= 0){
            ipCount = limit.permitsPerSecond();
        }
        String allKey = request.getRequestURI();
        String ipKey = request.getRequestURI() + IpAddressUtil.getIpAdrress(request);;

        RateLimiter allRateLimiter = null;
        //初始化总令牌请求次数
        if (!limitAllMap.containsKey(allKey)) {
            allRateLimiter = RateLimiter.create(limit.permitsPerSecond());
            limitAllMap.put(allKey, allRateLimiter);
            log.info("新建all令牌桶={}，容量={}",allRateLimiter,limit.permitsPerSecond());
        }

        RateLimiter ipRateLimiter = null;
        //初始化指定令牌请求次数
        if (!limitIpMap.containsKey(ipKey)) {
            ipRateLimiter = RateLimiter.create(ipCount);
            limitIpMap.put(ipKey, ipRateLimiter);
            log.info("新建ip令牌桶={}，容量={}",ipRateLimiter,ipCount);
        }
        allRateLimiter = limitAllMap.get(allKey);
        ipRateLimiter = limitIpMap.get(ipKey);

        boolean allAcquire = allRateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
        boolean ipAcquire = ipRateLimiter.tryAcquire(limit.timeout(), limit.timeunit());

        if (!allAcquire) {
            log.info("总令牌桶={}，获取令牌失败",allKey);
            return Result.fail(limit.msg());
        }
        if(!ipAcquire){
            log.info("ip牌桶={}，获取令牌失败",ipKey);
            return Result.fail(limit.msg());
        }
        return joinPoint.proceed();
    }


}
